use std::{fs, path::PathBuf};

use burn::{
    backend::{Autodiff, LibTorch, libtorch::LibTorchDevice},
    optim::AdamConfig,
};
use clap::Parser;
use hwdb::{
    data::{
        self,
        dataset::{HWDBDatasetConfig, HWDBHdf5DatasetConfig},
        utils::{read_classes_file, read_hdf5_classes_len, read_subdir},
    },
    model::basic::BasicModelConfig,
    train::{
        infer::infer,
        train::{HWDBDatasetConfType, TrainingConfig, train},
    },
};

use crate::arg_parse::{Config, Operator};

mod arg_parse;

fn main() {
    let config = Config::parse();
    match config.op {
        Operator::Transfer {
            raw_dir,
            images_dir,
        } => {
            if !fs::exists(&images_dir).map_or(false, |exist| exist) {
                fs::create_dir_all(&images_dir).expect("can't create images dir");
            }
            data::transfer::batch_transfer(&raw_dir, &images_dir).expect("Transfer failed");
        }
        Operator::Unzip {
            zip_dir,
            save_dir,
            smart_check,
            trn_list,
            tst_list,
            vld_list,
            trn_tag,
            tst_tag,
            vld_tag,
            hdf5,
            use_hdf5,
        } => {
            data::transfer::batch_unzip_and_transfer(
                &zip_dir,
                &save_dir,
                smart_check,
                trn_list,
                tst_list,
                vld_list,
                trn_tag,
                tst_tag,
                vld_tag,
                if use_hdf5 { Some(&hdf5) } else { None },
            )
            .expect("Unzip failed");
        }
        Operator::Preprocess {
            images_dir,
            classes_filename,
            label_filename,
            trn_tag,
            tst_tag,
            vld_tag,
        } => {
            let root_dir = PathBuf::from(&images_dir);
            assert!(root_dir.exists(), "images dir is not exist");

            let mut parts = Vec::new();

            let trn_dir = root_dir.join(&trn_tag);
            let tst_dir = root_dir.join(&tst_tag);
            let vld_dir = root_dir.join(&vld_tag);

            if trn_dir.exists() {
                parts.extend_from_slice(&read_subdir(trn_dir.as_path()).as_slice());
            }

            if tst_dir.exists() {
                parts.extend_from_slice(&read_subdir(tst_dir.as_path()).as_slice());
            }

            if vld_dir.exists() {
                parts.extend_from_slice(&read_subdir(vld_dir.as_path()).as_slice());
            }

            let classes_filepath = root_dir.join(classes_filename);
            data::preprocess::preprocess_images(
                images_dir,
                parts,
                classes_filepath
                    .to_str()
                    .expect("can't find classes file")
                    .to_string(),
                label_filename,
            )
            .expect("preprocess error");
        }
        Operator::Train {
            images_dir,
            log_dir,
            classes_filename,
            label_filename,
            num_epoches,
            batch_size,
            num_workers,
            seed,
            lr,
            image_size,
            trn_tag,
            tst_tag,
            hdf5,
            use_hdf5,
        } => {
            let root_dir = PathBuf::from(&images_dir);
            assert!(root_dir.exists(), "hdf5 dir is not exist");

            let dataset = if use_hdf5 {
                let hdf5_filepath = root_dir.join(&hdf5);
                assert!(hdf5_filepath.exists(), "hdf5 file is not exist");

                HWDBDatasetConfType::Hdf5(HWDBHdf5DatasetConfig::new(
                    hdf5_filepath.display().to_string(),
                ))
            } else {
                let trn_dirpath = root_dir.join(&trn_tag);
                assert!(
                    trn_dirpath.exists(),
                    "train dir is not exist, train tag: {}",
                    trn_tag
                );

                let tst_dirpath = root_dir.join(&tst_tag);
                assert!(
                    tst_dirpath.exists(),
                    "test dir is not exist, test tag: {}",
                    tst_tag
                );

                let trn_dir_vec = read_subdir(trn_dirpath.as_path());
                let tst_dir_vec = read_subdir(tst_dirpath.as_path());

                let classes_path = root_dir.join(&classes_filename);
                HWDBDatasetConfType::Folder(HWDBDatasetConfig::new(
                    trn_dir_vec,
                    tst_dir_vec,
                    classes_path.to_str().unwrap().to_string(),
                    label_filename,
                ))
            };

            let model = BasicModelConfig::new(if use_hdf5 {
                let root_dir = PathBuf::from(&images_dir);
                assert!(root_dir.exists(), "hdf5 dir is not exist");

                let hdf5_filepath = root_dir.join(&hdf5);
                assert!(hdf5_filepath.exists(), "hdf5 file is not exist");

                read_hdf5_classes_len(&hdf5_filepath).expect("can't read hdf5 classes len")
            } else {
                let classes_path = root_dir.join(&classes_filename);
                assert!(
                    classes_path.exists(),
                    "classes file name is not exist {}",
                    classes_path.display()
                );
                let classes =
                    read_classes_file(classes_path.to_str().expect("classes file path is invalid"))
                        .expect("read classes file failed");
                classes.len()
            });
            let optimizer = AdamConfig::new();
            let training_config = TrainingConfig::new(dataset, model, optimizer)
                .with_num_epoches(num_epoches)
                .with_batch_size(batch_size)
                .with_lr(lr)
                .with_num_workers(num_workers)
                .with_seed(seed)
                .with_image_size(image_size);
            type Backend = Autodiff<LibTorch>;
            println!(
                "check available CUDA, total size: {}, Using CUDA: {}",
                tch::Cuda::device_count(),
                tch::Cuda::is_available()
            );
            let device = if tch::Cuda::is_available() {
                LibTorchDevice::Cuda(0)
            } else {
                Default::default()
            };
            train::<Backend>(&log_dir, training_config, device);
        }
        Operator::Infer {
            log_dir,
            images_path,
            h5_filepath,
            intermediate_image_save_path,
        } => {
            type Backend = LibTorch;
            println!(
                "check available CUDA, total size: {}, Using CUDA: {}",
                tch::Cuda::device_count(),
                tch::Cuda::is_available()
            );
            let device = if tch::Cuda::is_available() {
                LibTorchDevice::Cuda(0)
            } else {
                Default::default()
            };
            infer::<Backend>(
                &log_dir,
                device,
                images_path,
                h5_filepath,
                intermediate_image_save_path,
            );
        }
    }
}
